# 导入所需的库
from pathlib import Path  # 用于处理文件路径
import numpy as np  # 用于数值计算
import torch  # PyTorch深度学习框架
import torch.nn.functional as F  # PyTorch函数式接口
from tqdm import tqdm  # 用于显示进度条

from dataset import save_emb  # 从dataset模块导入save_emb函数


class FlashMultiHeadAttention(torch.nn.Module):
    """
    实现Flash Multi-Head Attention机制的类
    支持PyTorch 2.0+的内置Flash Attention和降级的标准Attention
    """
    def __init__(self, hidden_units, num_heads, dropout_rate):
        # 调用父类构造函数
        super(FlashMultiHeadAttention, self).__init__()

        # 存储模型参数
        self.hidden_units = hidden_units  # 隐藏单元数
        self.num_heads = num_heads  # 注意力头数
        self.head_dim = hidden_units // num_heads  # 每个头的维度
        self.dropout_rate = dropout_rate  # dropout率

        # 确保隐藏单元数能被头数整除
        assert hidden_units % num_heads == 0, "hidden_units must be divisible by num_heads"

        # 定义线性变换层
        self.q_linear = torch.nn.Linear(hidden_units, hidden_units)  # Query线性变换
        self.k_linear = torch.nn.Linear(hidden_units, hidden_units)  # Key线性变换
        self.v_linear = torch.nn.Linear(hidden_units, hidden_units)  # Value线性变换
        self.out_linear = torch.nn.Linear(hidden_units, hidden_units)  # 输出线性变换

    def forward(self, query, key, value, attn_mask=None):
        # 获取输入张量的形状
        batch_size, seq_len, _ = query.size()

        # 计算Q, K, V
        Q = self.q_linear(query)  # 查询向量
        K = self.k_linear(key)    # 键向量
        V = self.v_linear(value)  # 值向量

        # reshape为multi-head格式: (batch_size, num_heads, seq_len, head_dim)
        Q = Q.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        K = K.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        V = V.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)

        # 检查是否有PyTorch 2.0+的内置Flash Attention
        if hasattr(F, 'scaled_dot_product_attention'):
            # PyTorch 2.0+ 使用内置的Flash Attention
            attn_output = F.scaled_dot_product_attention(
                Q, K, V, dropout_p=self.dropout_rate if self.training else 0.0, attn_mask=attn_mask.unsqueeze(1)
            )
        else:
            # 降级到标准注意力机制
            scale = (self.head_dim) ** -0.5  # 缩放因子
            scores = torch.matmul(Q, K.transpose(-2, -1)) * scale  # 计算注意力分数

            # 如果有注意力掩码，则应用掩码
            if attn_mask is not None:
                scores.masked_fill_(attn_mask.unsqueeze(1).logical_not(), float('-inf'))

            # 计算注意力权重并应用dropout
            attn_weights = F.softmax(scores, dim=-1)
            attn_weights = F.dropout(attn_weights, p=self.dropout_rate, training=self.training)
            # 计算注意力输出
            attn_output = torch.matmul(attn_weights, V)

        # reshape回原来的格式: (batch_size, seq_len, hidden_units)
        attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, seq_len, self.hidden_units)

        # 最终的线性变换
        output = self.out_linear(attn_output)

        # 返回输出和None(兼容性考虑)
        return output, None


class PointWiseFeedForward(torch.nn.Module):
    """
    实现逐点前馈网络的类
    使用1D卷积实现位置无关的前馈变换
    """
    def __init__(self, hidden_units, dropout_rate):
        # 调用父类构造函数
        super(PointWiseFeedForward, self).__init__()

        # 定义前馈网络的各层
        self.conv1 = torch.nn.Conv1d(hidden_units, hidden_units, kernel_size=1)  # 第一个1D卷积层
        self.dropout1 = torch.nn.Dropout(p=dropout_rate)  # 第一个dropout层
        self.relu = torch.nn.ReLU()  # ReLU激活函数
        self.conv2 = torch.nn.Conv1d(hidden_units, hidden_units, kernel_size=1)  # 第二个1D卷积层
        self.dropout2 = torch.nn.Dropout(p=dropout_rate)  # 第二个dropout层

    def forward(self, inputs):
        # 前馈网络的前向传播
        # 调整输入维度以适应Conv1D要求: (N, C, Length)
        outputs = self.dropout2(self.conv2(self.relu(self.dropout1(self.conv1(inputs.transpose(-1, -2))))))
        # 转换回原始维度: (N, Length, C)
        outputs = outputs.transpose(-1, -2)
        return outputs


class BaselineModel(torch.nn.Module):
    """
    基线模型类，实现了一个基于Transformer的序列推荐模型
    Args:
        user_num: 用户数量
        item_num: 物品数量
        feat_statistics: 特征统计信息，key为特征ID，value为特征数量
        feat_types: 各个特征的特征类型，key为特征类型名称，value为包含的特征ID列表，包括user和item的sparse, array, emb, continual类型
        args: 全局参数

    Attributes:
        user_num: 用户数量
        item_num: 物品数量
        dev: 设备
        norm_first: 是否先归一化
        maxlen: 序列最大长度
        item_emb: Item Embedding Table
        user_emb: User Embedding Table
        sparse_emb: 稀疏特征Embedding Table
        emb_transform: 多模态特征的线性变换
        userdnn: 用户特征拼接后经过的全连接层
        itemdnn: 物品特征拼接后经过的全连接层
    """

    def __init__(self, user_num, item_num, feat_statistics, feat_types, args):
        # 调用父类构造函数
        super(BaselineModel, self).__init__()

        # 存储基本参数
        self.user_num = user_num  # 用户数量
        self.item_num = item_num  # 物品数量
        self.dev = args.device    # 设备(CPU/GPU)
        self.norm_first = args.norm_first  # 是否先进行LayerNorm
        self.maxlen = args.maxlen  # 序列最大长度
        # TODO: loss += args.l2_emb for regularizing embedding vectors during training
        # https://stackoverflow.com/questions/42704283/adding-l1-l2-regularization-in-pytorch

        # 定义基础嵌入层
        self.item_emb = torch.nn.Embedding(self.item_num + 1, args.hidden_units, padding_idx=0)  # 物品嵌入表
        self.user_emb = torch.nn.Embedding(self.user_num + 1, args.hidden_units, padding_idx=0)  # 用户嵌入表
        self.pos_emb = torch.nn.Embedding(2 * args.maxlen + 1, args.hidden_units, padding_idx=0)  # 位置嵌入表
        self.emb_dropout = torch.nn.Dropout(p=args.dropout_rate)  # 嵌入层dropout

        # 定义特征嵌入和变换模块
        self.sparse_emb = torch.nn.ModuleDict()  # 稀疏特征嵌入字典
        self.emb_transform = torch.nn.ModuleDict()  # 多模态特征变换字典

        # 定义Transformer层组件
        self.attention_layernorms = torch.nn.ModuleList()  # 注意力层归一化列表
        self.attention_layers = torch.nn.ModuleList()      # 注意力层列表
        self.forward_layernorms = torch.nn.ModuleList()    # 前馈层归一化列表
        self.forward_layers = torch.nn.ModuleList()        # 前馈层列表

        # 初始化特征信息
        self._init_feat_info(feat_statistics, feat_types)

        # 计算用户和物品特征维度
        userdim = args.hidden_units * (len(self.USER_SPARSE_FEAT) + 1 + len(self.USER_ARRAY_FEAT)) + len(
            self.USER_CONTINUAL_FEAT
        )
        itemdim = (
            args.hidden_units * (len(self.ITEM_SPARSE_FEAT) + 1 + len(self.ITEM_ARRAY_FEAT))
            + len(self.ITEM_CONTINUAL_FEAT)
            + args.hidden_units * len(self.ITEM_EMB_FEAT)
        )

        # 定义用户和物品特征的全连接层
        self.userdnn = torch.nn.Linear(userdim, args.hidden_units)  # 用户特征全连接层
        self.itemdnn = torch.nn.Linear(itemdim, args.hidden_units)  # 物品特征全连接层

        # 最后一层归一化
        self.last_layernorm = torch.nn.LayerNorm(args.hidden_units, eps=1e-8)

        # 创建Transformer块
        for _ in range(args.num_blocks):
            # 添加注意力层归一化
            new_attn_layernorm = torch.nn.LayerNorm(args.hidden_units, eps=1e-8)
            self.attention_layernorms.append(new_attn_layernorm)

            # 添加注意力层(使用FlashAttention)
            new_attn_layer = FlashMultiHeadAttention(
                args.hidden_units, args.num_heads, args.dropout_rate
            )  # 优化：用FlashAttention替代标准Attention
            self.attention_layers.append(new_attn_layer)

            # 添加前馈层归一化
            new_fwd_layernorm = torch.nn.LayerNorm(args.hidden_units, eps=1e-8)
            self.forward_layernorms.append(new_fwd_layernorm)

            # 添加前馈层
            new_fwd_layer = PointWiseFeedForward(args.hidden_units, args.dropout_rate)
            self.forward_layers.append(new_fwd_layer)

        # 为各种稀疏特征创建嵌入表
        for k in self.USER_SPARSE_FEAT:
            self.sparse_emb[k] = torch.nn.Embedding(self.USER_SPARSE_FEAT[k] + 1, args.hidden_units, padding_idx=0)
        for k in self.ITEM_SPARSE_FEAT:
            self.sparse_emb[k] = torch.nn.Embedding(self.ITEM_SPARSE_FEAT[k] + 1, args.hidden_units, padding_idx=0)
        for k in self.ITEM_ARRAY_FEAT:
            self.sparse_emb[k] = torch.nn.Embedding(self.ITEM_ARRAY_FEAT[k] + 1, args.hidden_units, padding_idx=0)
        for k in self.USER_ARRAY_FEAT:
            self.sparse_emb[k] = torch.nn.Embedding(self.USER_ARRAY_FEAT[k] + 1, args.hidden_units, padding_idx=0)
        # 为多模态特征创建线性变换层
        for k in self.ITEM_EMB_FEAT:
            self.emb_transform[k] = torch.nn.Linear(self.ITEM_EMB_FEAT[k], args.hidden_units)

    def _init_feat_info(self, feat_statistics, feat_types):
        """
        将特征统计信息（特征数量）按特征类型分组产生不同的字典，方便声明稀疏特征的Embedding Table

        Args:
            feat_statistics: 特征统计信息，key为特征ID，value为特征数量
            feat_types: 各个特征的特征类型，key为特征类型名称，value为包含的特征ID列表，包括user和item的sparse, array, emb, continual类型
        """
        # 按类型分组特征
        self.USER_SPARSE_FEAT = {k: feat_statistics[k] for k in feat_types['user_sparse']}      # 用户稀疏特征
        self.USER_CONTINUAL_FEAT = feat_types['user_continual']                                 # 用户连续特征
        self.ITEM_SPARSE_FEAT = {k: feat_statistics[k] for k in feat_types['item_sparse']}      # 物品稀疏特征
        self.ITEM_CONTINUAL_FEAT = feat_types['item_continual']                                 # 物品连续特征
        self.USER_ARRAY_FEAT = {k: feat_statistics[k] for k in feat_types['user_array']}        # 用户数组特征
        self.ITEM_ARRAY_FEAT = {k: feat_statistics[k] for k in feat_types['item_array']}        # 物品数组特征
        
        # 多模态特征维度字典
        EMB_SHAPE_DICT = {"81": 32, "82": 1024, "83": 3584, "84": 4096, "85": 3584, "86": 3584}
        self.ITEM_EMB_FEAT = {k: EMB_SHAPE_DICT[k] for k in feat_types['item_emb']}  # 记录的是不同多模态特征的维度

    def feat2tensor(self, seq_feature, k):
        """
        将序列特征转换为张量

        Args:
            seq_feature: 序列特征list，每个元素为当前时刻的特征字典，形状为 [batch_size, maxlen]
            k: 特征ID

        Returns:
            batch_data: 特征值的tensor，形状为 [batch_size, maxlen, max_array_len(if array)]
        """
        batch_size = len(seq_feature)  # 批次大小

        # 判断特征是否为数组类型
        if k in self.ITEM_ARRAY_FEAT or k in self.USER_ARRAY_FEAT:
            # 如果特征是Array类型，需要先对array进行padding，然后转换为tensor
            max_array_len = 0  # 最大数组长度
            max_seq_len = 0    # 最大序列长度

            # 计算最大数组长度和最大序列长度
            for i in range(batch_size):
                seq_data = [item[k] for item in seq_feature[i]]  # 获取序列数据
                max_seq_len = max(max_seq_len, len(seq_data))
                max_array_len = max(max_array_len, max(len(item_data) for item_data in seq_data))

            # 创建零填充的数组
            batch_data = np.zeros((batch_size, max_seq_len, max_array_len), dtype=np.int64)
            # 填充实际数据
            for i in range(batch_size):
                seq_data = [item[k] for item in seq_feature[i]]
                for j, item_data in enumerate(seq_data):
                    actual_len = min(len(item_data), max_array_len)
                    batch_data[i, j, :actual_len] = item_data[:actual_len]

            # 转换为张量并移动到指定设备
            return torch.from_numpy(batch_data).to(self.dev)
        else:
            # 如果特征是Sparse类型，直接转换为tensor
            max_seq_len = max(len(seq_feature[i]) for i in range(batch_size))
            batch_data = np.zeros((batch_size, max_seq_len), dtype=np.int64)

            # 填充数据
            for i in range(batch_size):
                seq_data = [item[k] for item in seq_feature[i]]
                batch_data[i] = seq_data

            # 转换为张量并移动到指定设备
            return torch.from_numpy(batch_data).to(self.dev)

    def feat2emb(self, seq, feature_array, mask=None, include_user=False):
        """
        将特征转换为嵌入向量

        Args:
            seq: 序列ID
            feature_array: 特征list，每个元素为当前时刻的特征字典
            mask: 掩码，1表示item，2表示user
            include_user: 是否处理用户特征，在两种情况下不打开：1) 训练时在转换正负样本的特征时（因为正负样本都是item）;2) 生成候选库item embedding时。

        Returns:
            seqs_emb: 序列特征的Embedding
        """
        seq = seq.to(self.dev)  # 将序列移动到指定设备
        
        # 预计算嵌入
        if include_user:
            # 如果包含用户特征
            user_mask = (mask == 2).to(self.dev)  # 用户掩码
            item_mask = (mask == 1).to(self.dev)  # 物品掩码
            user_embedding = self.user_emb(user_mask * seq)  # 用户嵌入
            item_embedding = self.item_emb(item_mask * seq)  # 物品嵌入
            item_feat_list = [item_embedding]  # 物品特征列表
            user_feat_list = [user_embedding]  # 用户特征列表
        else:
            # 仅处理物品特征
            item_embedding = self.item_emb(seq)
            item_feat_list = [item_embedding]

        # 批处理所有特征类型
        all_feat_types = [
            (self.ITEM_SPARSE_FEAT, 'item_sparse', item_feat_list),   # 物品稀疏特征
            (self.ITEM_ARRAY_FEAT, 'item_array', item_feat_list),     # 物品数组特征
            (self.ITEM_CONTINUAL_FEAT, 'item_continual', item_feat_list),  # 物品连续特征
        ]

        # 如果包含用户特征，则添加用户特征类型
        if include_user:
            all_feat_types.extend(
                [
                    (self.USER_SPARSE_FEAT, 'user_sparse', user_feat_list),    # 用户稀疏特征
                    (self.USER_ARRAY_FEAT, 'user_array', user_feat_list),      # 用户数组特征
                    (self.USER_CONTINUAL_FEAT, 'user_continual', user_feat_list),  # 用户连续特征
                ]
            )

        # 批处理每种特征类型
        for feat_dict, feat_type, feat_list in all_feat_types:
            if not feat_dict:  # 如果特征字典为空则跳过
                continue

            for k in feat_dict:
                # 将特征转换为张量
                tensor_feature = self.feat2tensor(feature_array, k)

                # 根据特征类型应用不同的处理方式
                if feat_type.endswith('sparse'):
                    # 稀疏特征直接查嵌入表
                    feat_list.append(self.sparse_emb[k](tensor_feature))
                elif feat_type.endswith('array'):
                    # 数组特征查嵌入表后求和
                    feat_list.append(self.sparse_emb[k](tensor_feature).sum(2))
                elif feat_type.endswith('continual'):
                    # 连续特征增加维度
                    feat_list.append(tensor_feature.unsqueeze(2))

        # 处理多模态特征
        for k in self.ITEM_EMB_FEAT:
            # 收集所有数据到numpy数组，然后批量转换
            batch_size = len(feature_array)
            emb_dim = self.ITEM_EMB_FEAT[k]  # 嵌入维度
            seq_len = len(feature_array[0])   # 序列长度

            # 预分配张量
            batch_emb_data = np.zeros((batch_size, seq_len, emb_dim), dtype=np.float32)

            # 填充数据
            for i, seq in enumerate(feature_array):
                for j, item in enumerate(seq):
                    if k in item:
                        batch_emb_data[i, j] = item[k]

            # 批量转换并移动到GPU
            tensor_feature = torch.from_numpy(batch_emb_data).to(self.dev)
            item_feat_list.append(self.emb_transform[k](tensor_feature))

        # 合并特征
        all_item_emb = torch.cat(item_feat_list, dim=2)  # 拼接物品特征
        all_item_emb = torch.relu(self.itemdnn(all_item_emb))  # 通过物品全连接层
        if include_user:
            all_user_emb = torch.cat(user_feat_list, dim=2)   # 拼接用户特征
            all_user_emb = torch.relu(self.userdnn(all_user_emb))  # 通过用户全连接层
            seqs_emb = all_item_emb + all_user_emb  # 合并物品和用户特征
        else:
            seqs_emb = all_item_emb  # 仅使用物品特征
        return seqs_emb

    def log2feats(self, log_seqs, mask, seq_feature):
        """
        将日志序列转换为特征表示

        Args:
            log_seqs: 序列ID
            mask: token类型掩码，1表示item token，2表示user token
            seq_feature: 序列特征list，每个元素为当前时刻的特征字典

        Returns:
            seqs_emb: 序列的Embedding，形状为 [batch_size, maxlen, hidden_units]
        """
        batch_size = log_seqs.shape[0]  # 批次大小
        maxlen = log_seqs.shape[1]      # 最大序列长度
        
        # 将特征转换为嵌入向量
        seqs = self.feat2emb(log_seqs, seq_feature, mask=mask, include_user=True)
        seqs *= self.item_emb.embedding_dim**0.5  # 缩放嵌入向量
        
        # 计算位置编码
        poss = torch.arange(1, maxlen + 1, device=self.dev).unsqueeze(0).expand(batch_size, -1).clone()
        poss *= log_seqs != 0  # 对padding位置应用掩码
        seqs += self.pos_emb(poss)  # 添加位置嵌入
        seqs = self.emb_dropout(seqs)  # 应用dropout

        # 计算序列长度和注意力掩码
        maxlen = seqs.shape[1]
        ones_matrix = torch.ones((maxlen, maxlen), dtype=torch.bool, device=self.dev)
        attention_mask_tril = torch.tril(ones_matrix)  # 下三角矩阵(防止未来信息泄露)
        attention_mask_pad = (mask != 0).to(self.dev)  # padding掩码
        attention_mask = attention_mask_tril.unsqueeze(0) & attention_mask_pad.unsqueeze(1)  # 组合掩码

        # 通过Transformer层
        for i in range(len(self.attention_layers)):
            if self.norm_first:
                # 先归一化再处理
                x = self.attention_layernorms[i](seqs)
                mha_outputs, _ = self.attention_layers[i](x, x, x, attn_mask=attention_mask)
                seqs = seqs + mha_outputs  # 残差连接
                seqs = seqs + self.forward_layers[i](self.forward_layernorms[i](seqs))  # 前馈网络和残差连接
            else:
                # 标准处理方式
                mha_outputs, _ = self.attention_layers[i](seqs, seqs, seqs, attn_mask=attention_mask)
                seqs = self.attention_layernorms[i](seqs + mha_outputs)  # 注意力和残差连接
                seqs = self.forward_layernorms[i](seqs + self.forward_layers[i](seqs))  # 前馈和残差连接

        # 最后一层归一化
        log_feats = self.last_layernorm(seqs)

        return log_feats

    def forward(
        self, user_item, pos_seqs, neg_seqs, mask, next_mask, next_action_type, seq_feature, pos_feature, neg_feature
    ):
        """
        训练时调用，计算正负样本的logits

        Args:
            user_item: 用户序列ID
            pos_seqs: 正样本序列ID
            neg_seqs: 负样本序列ID
            mask: token类型掩码，1表示item token，2表示user token
            next_mask: 下一个token类型掩码，1表示item token，2表示user token
            next_action_type: 下一个token动作类型，0表示曝光，1表示点击
            seq_feature: 序列特征list，每个元素为当前时刻的特征字典
            pos_feature: 正样本特征list，每个元素为当前时刻的特征字典
            neg_feature: 负样本特征list，每个元素为当前时刻的特征字典

        Returns:
            pos_logits: 正样本logits，形状为 [batch_size, maxlen]
            neg_logits: 负样本logits，形状为 [batch_size, maxlen]
        """
        # 获取序列特征表示
        log_feats = self.log2feats(user_item, mask, seq_feature)
        loss_mask = (next_mask == 1).to(self.dev)  # 损失掩码(仅考虑物品token)

        # 获取正负样本嵌入
        pos_embs = self.feat2emb(pos_seqs, pos_feature, include_user=False)
        neg_embs = self.feat2emb(neg_seqs, neg_feature, include_user=False)

        # 计算正负样本logits
        pos_logits = (log_feats * pos_embs).sum(dim=-1)  # 点积相似度
        neg_logits = (log_feats * neg_embs).sum(dim=-1)
        pos_logits = pos_logits * loss_mask  # 应用损失掩码
        neg_logits = neg_logits * loss_mask

        return pos_logits, neg_logits

    def predict(self, log_seqs, seq_feature, mask):
        """
        计算用户序列的表征
        Args:
            log_seqs: 用户序列ID
            seq_feature: 序列特征list，每个元素为当前时刻的特征字典
            mask: token类型掩码，1表示item token，2表示user token
        Returns:
            final_feat: 用户序列的表征，形状为 [batch_size, hidden_units]
        """
        # 获取序列特征表示
        log_feats = self.log2feats(log_seqs, mask, seq_feature)

        # 取最后一个时间步的特征作为用户表征
        final_feat = log_feats[:, -1, :]

        return final_feat

    def save_item_emb(self, item_ids, retrieval_ids, feat_dict, save_path, batch_size=1024):
        """
        生成候选库item embedding，用于检索

        Args:
            item_ids: 候选item ID（re-id形式）
            retrieval_ids: 候选item ID（检索ID，从0开始编号，检索脚本使用）
            feat_dict: 训练集所有item特征字典，key为特征ID，value为特征值
            save_path: 保存路径
            batch_size: 批次大小
        """
        all_embs = []  # 存储所有嵌入向量

        # 分批处理物品ID
        for start_idx in tqdm(range(0, len(item_ids), batch_size), desc="Saving item embeddings"):
            end_idx = min(start_idx + batch_size, len(item_ids))

            # 获取当前批次的物品ID
            item_seq = torch.tensor(item_ids[start_idx:end_idx], device=self.dev).unsqueeze(0)
            batch_feat = []
            # 收集当前批次的特征
            for i in range(start_idx, end_idx):
                batch_feat.append(feat_dict[i])

            batch_feat = np.array(batch_feat, dtype=object)

            # 计算嵌入向量
            batch_emb = self.feat2emb(item_seq, [batch_feat], include_user=False).squeeze(0)

            # 存储嵌入向量
            all_embs.append(batch_emb.detach().cpu().numpy().astype(np.float32))

        # 合并所有批次的结果并保存
        final_ids = np.array(retrieval_ids, dtype=np.uint64).reshape(-1, 1)
        final_embs = np.concatenate(all_embs, axis=0)
        save_emb(final_embs, Path(save_path, 'embedding.fbin'))  # 保存嵌入向量
        save_emb(final_ids, Path(save_path, 'id.u64bin'))       # 保存ID